#!/usr/bin/python
# -*- coding: utf-8 -*-
import pandas as pd
from pandas import DataFrame
from numpy import nan
from enum import Enum

from core.index_base import IndexBase


class OBVType(Enum):
    accumulation = 1 # 积累
    moving = 2 # 移动
    revision = 3 # 修正


class OBV(IndexBase):
    """
    On Balance Volume 能量潮
    """
    def __init__(self, period=9, type=OBVType.accumulation, key='close'):
        super(OBV, self).__init__()
        self.period = period
        self.type = type
        self.key = key

    def compute(self, data:DataFrame):
        pre_value = 0
        if self.type == OBVType.moving or self.type == OBVType.accumulation:
            obv_key = 'obv'
            diff_key = 'obv_diff'
            self.added_keys.extend([obv_key, diff_key])

            data[obv_key] = nan
            data[diff_key] = nan

            data[diff_key] = data[self.key].diff()
            pre_value = 0
            flag = 1
            for i in [str(date)[:10] for date in data.index.values]:
                pre_value = pre_value if not pd.isnull(pre_value) else 0
                flag = data.loc[i, diff_key]/abs(data.loc[i, diff_key])
                flag = flag if not pd.isnull(flag) else 1
                pre_value = flag * data.loc[i,'vol'] + pre_value
                data.loc[i,obv_key] = pre_value
            if self.type == OBVType.moving:
                obv_mov_key = 'obv-mov' + self.period
                self.added_keys.extend([obv_mov_key])

                data[obv_mov_key] = nan
                data[obv_mov_key] = data[obv_key].rolling(self.period).mean()
        elif self.type == OBVType.revision:
            obv_rev_key = 'obv_rev'
            self.added_keys.append(obv_rev_key)

            data[obv_rev_key] = (((data['close'] - data['low']) - (data['high'] - data['close']))/(data['high'] - data['low'])) * data['vol']
            for i in [str(date)[:10] for date in data.index.values]:
                pre_value = pre_value if not pd.isnull(pre_value) else 0
                pre_value = pre_value + data.loc[i,obv_rev_key]
                data.loc[i, obv_rev_key] = pre_value


if __name__=="__main__":
    from domain.transaction_data.repository import transaction_data_repository

    OBV(type=OBVType.revision).compute(transaction_data_repository.get_contain_latest_data_by_tushare('002044.SZ'))